#!/bin/bash

# 系统软件源定制

# 获取操作系统名称
system=$(hostnamectl | grep -i operating | awk '{print $3}')

# Ubuntu系统源修改函数
chubantu_repo() {
    archive=$(grep "URIs:" /etc/apt/sources.list.d/ubuntu.sources | awk 'NR==1 {print $2}')
    security=$(grep "URIs:" /etc/apt/sources.list.d/ubuntu.sources | awk 'NR==2 {print $2}')

    if [ -z "$archive" ]; then
        echo "错误：未找到archive源的URL" >&2
        exit 1
    else
        # 修改源配置文件
        sed -i.bak -e "s|$archive|https://mirrors.aliyun.com/ubuntu|g" \
                   -e "s|$security|https://mirrors.aliyun.com/ubuntu|g" \
                   /etc/apt/sources.list.d/ubuntu.sources
    fi
    
    echo "修改后的配置："
    grep "URIs:" /etc/apt/sources.list.d/ubuntu.sources
}

# Rocky系统源修改函数
chrocky_repo() {
    # 检查配置文件是否存在
    if ! ls /etc/yum.repos.d/rocky*.repo >/dev/null 2>&1; then
        echo "错误：未找到Rocky Linux配置文件" >&2
        exit 1
    else
        # 修改源配置文件
        sed -i.bak -e 's|^mirrorlist=|#mirrorlist=|g' \
                   -e 's|^#baseurl=http://dl.rockylinux.org/$contentdir|baseurl=https://mirrors.aliyun.com/rockylinux|g' \
                   /etc/yum.repos.d/rocky*.repo
    fi
}

main() {
    if [ "$system" = "Ubuntu" ]; then
        chubantu_repo
        echo "正在更新软件包索引..."
        if apt update; then
            echo "成功更新软件包索引"
        else
            echo "更新失败，正在恢复原始配置"
            cp /etc/apt/sources.list.d/ubuntu.sources.bak /etc/apt/sources.list.d/ubuntu.sources
            apt update
            exit 1
        fi
    elif [ "$system" = "Rocky" ]; then
        chrocky_repo
        echo "正在更新软件包缓存..."
        if yum makecache; then
            echo "成功更新软件包缓存"
        else
            echo "更新失败，正在恢复原始配置"
            for file in /etc/yum.repos.d/*.bak; do
                mv "$file" "${file%.bak}"
            done
            yum makecache
            exit 1
        fi
    else
        echo "错误：不支持的操作系统 $system" >&2
        exit 1
    fi
}

main
